from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
import os
import base64

# AES 密钥生成 (16, 24, 或 32 字节长)
SECRET_KEY = b'secretpassword12'  # 示例密钥，16字节


# AES 加密函数
def encrypt_aes(data: str) -> str:
    # 生成随机的16字节初始化向量 (IV)
    iv = os.urandom(16)

    # 创建 AES CBC 加密器
    cipher = Cipher(algorithms.AES(SECRET_KEY), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()

    # 使用 PKCS7 填充数据，使其成为 16 字节的倍数
    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(data.encode('utf-8')) + padder.finalize()

    # 加密数据
    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()

    # 将 IV 和加密后的数据组合在一起，进行 Base64 编码返回
    return base64.b64encode(iv + encrypted_data).decode('utf-8')


# AES 解密函数
def decrypt_aes(encrypted_data: str) -> str:
    # Base64 解码加密后的数据
    encrypted_data = base64.b64decode(encrypted_data)

    # 提取 IV 和加密后的数据
    iv = encrypted_data[:16]  # 前16字节是 IV
    encrypted_data = encrypted_data[16:]  # 剩下的才是加密的数据

    # 创建 AES CBC 解密器
    cipher = Cipher(algorithms.AES(SECRET_KEY), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()

    # 解密数据
    padded_data = decryptor.update(encrypted_data) + decryptor.finalize()

    # 使用 PKCS7 去除填充
    unpadder = padding.PKCS7(128).unpadder()
    data = unpadder.update(padded_data) + unpadder.finalize()

    return data.decode('utf-8')


# 示例用法
if __name__ == "__main__":
    data = "Hello, this is a secret message!"

    # 加密
    encrypted_data = encrypt_aes(data)
    print(f"Encrypted Data: {encrypted_data}")

    # 解密
    decrypted_data = decrypt_aes(encrypted_data)
    print(f"Decrypted Data: {decrypted_data}")
